//
//  OLED25.swift
//  videoPlayerSwift
//
//  Created by zhe wu on 2025/11/21.
//  Copyright © 2025 xuwanliang. All rights reserved.
//

import MetalKit
import Foundation
import SwiftUI

class OLED25: NSObject {
    var device: MTLDevice?
    var commandQueue: MTLCommandQueue! = nil
    var computePipeline: MTLComputePipelineState!
    var colorSamplerState: MTLSamplerState! = nil
    
    var startTime = Date()
    var width:Int = 0
    var height:Int = 0
    
    required init(device: MTLDevice, commandQueue: MTLCommandQueue, width: Int, height: Int){
        super.init()
        self.device = device
        self.commandQueue = commandQueue
        self.width = width
        self.height = height
        
        let library = device.makeDefaultLibrary()!
        let updateFunction = library.makeFunction(name: "oled25")!
        do{
            self.computePipeline = try device.makeComputePipelineState(function: updateFunction)
            print("OLED25 inited!")
        }catch let error as NSError{
            fatalError("\(error)")
        }
        
        let samplerDescriptor = MTLSamplerDescriptor()
        samplerDescriptor.minFilter = .linear
        samplerDescriptor.magFilter = .linear
        samplerDescriptor.sAddressMode = .repeat
        samplerDescriptor.tAddressMode = .repeat
        self.colorSamplerState = self.device?.makeSamplerState(descriptor: samplerDescriptor)
    }
    
    

    
    public func render(drawable: CAMetalDrawable!, rpd: MTLRenderPassDescriptor!, texture: MTLTexture!){
        
        guard let commandBuffer = commandQueue.makeCommandBuffer(),
              let computeEncoder = commandBuffer.makeComputeCommandEncoder() else {
            return
        }
        
        // Set the compute pipeline state
        computeEncoder.setComputePipelineState(computePipeline)
        
        // Get the output texture from the drawable
        let outTexture = drawable.texture
        
        // Set the output texture
        computeEncoder.setTexture(outTexture, index: 0)
        computeEncoder.setTexture(texture, index: 1)
        computeEncoder.setSamplerState(self.colorSamplerState, index: 0)
        // Feed in any parameters you need in the shader (using just time here)
        let currentTime = Date().timeIntervalSince(startTime)
        var params = ExampleShaderParams(time: Float(currentTime))
        computeEncoder.setBytes(&params, length: MemoryLayout<ExampleShaderParams>.size, index: 0)
        
        // Calculate thread group sizes
        let w = computePipeline.threadExecutionWidth
        let h = computePipeline.maxTotalThreadsPerThreadgroup / w
        let threadsPerThreadgroup = MTLSizeMake(w, h, 1)
        
        // Calculate thread groups based on texture size
        let threadgroupsX = (outTexture.width + w - 1) / w
        let threadgroupsY = (outTexture.height + h - 1) / h
        let threadgroups = MTLSizeMake(threadgroupsX, threadgroupsY, 1)
        
        // Dispatch the compute shader
        computeEncoder.dispatchThreadgroups(threadgroups, threadsPerThreadgroup: threadsPerThreadgroup)
        
        // End encoding and commit
        computeEncoder.endEncoding()
        commandBuffer.present(drawable)
        commandBuffer.commit()

    }
}
